/*
 * nbd.c
 *
 * Copyright (C) 2019 Alex A. <coderain@sdf.org>
 *
 * This program is free software: you can redistribute it and/or modify
 * it under the terms of the GNU Affero General Public License as
 * published by the Free Software Foundation, either version 3 of the
 * License, or (at your option) any later version.
 *
 * This program is distributed in the hope that it will be useful,
 * but WITHOUT ANY WARRANTY; without even the implied warranty of
 * MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
 * GNU Affero General Public License for more details.
 *
 * You should have received a copy of the GNU Affero General Public License
 * along with this program.  If not, see <http://www.gnu.org/licenses/>.
 */

#include <stdlib.h>
#include <stdio.h>
#include <endian.h>
#include <assert.h>
#include <pthread.h>
#include <unistd.h>
#include <errno.h>
#include <sys/socket.h>
#include <sys/un.h>
#include <netinet/in.h>
#include <arpa/inet.h>
#include <netdb.h>
#include "nbd.h"

#ifdef DEBUG
#define TRACE(...) fprintf(stderr, __VA_ARGS__)
#else
#define TRACE(...)
#endif

static inline bool read_uint16(int sockfd, uint16_t *value)
{
    if (read(sockfd, value, sizeof(*value)) != sizeof(*value)) return false;
    *value = be16toh(*value);
    return true;
}

static inline bool read_uint32(int sockfd, uint32_t *value)
{
    if (read(sockfd, value, sizeof(*value)) != sizeof(*value)) return false;
    *value = be32toh(*value);
    return true;
}

static inline bool read_uint64(int sockfd, uint64_t *value)
{
    if (read(sockfd, value, sizeof(*value)) != sizeof(*value)) return false;
    *value = be64toh(*value);
    return true;
}

static size_t tx_hash(const void *tx)
{
    return ((const nbd_transaction_t*)tx)->id % BUCKETS;
}

static int tx_compare(const void *p1, const void *p2)
{
    const nbd_transaction_t *a = p1, *b = p2;
    if (a->id < b->id) return -1;
    else if (a->id > b->id) return 1;
    else return 0;
}

static void *process_messages(void *param)
{
    uint32_t header;
    nbd_conn_t *conn = param;
    TRACE("Message dispatch thread initialized\n");

    while (read_uint32(conn->sockfd, &header))
    {
        if (header != NBD_SIMPLE_REPLY_MAGIC) continue;

        uint32_t error;
        uint64_t txid;
        if (!read_uint32(conn->sockfd, &error)) break;
        if (!read_uint64(conn->sockfd, &txid)) break;

        pthread_mutex_lock(&conn->mutex);
        nbd_transaction_t template = { .id = txid };
        nbd_transaction_t *tx = hash_table_search(&conn->tx_hash_table, &template, 0);

        if (!tx)
        {
            TRACE("WARNING: Received a reply for a non-existent or discarded transaction %llu\n", txid);
            goto next;
        }

        TRACE("[%llu] Response (status: %u)\n", txid, error);

        if (tx->command == NBD_CMD_READ && error == 0)
        {
            ssize_t bytes_read = read(conn->sockfd, tx->reply, tx->length);
            if (bytes_read >= 0) tx->length = bytes_read;
            else error = errno;
        }

        tx->error_code = error;
        tx->complete = 1;
        pthread_cond_signal(&tx->cond);

    next:
        pthread_mutex_unlock(&conn->mutex);
    }

    TRACE("Lost connection to the server\n");
    pthread_mutex_lock(&conn->mutex);

    void **ptr = NULL;
    while (ptr = hash_table_enum(&conn->tx_hash_table, ptr))
    {
        nbd_transaction_t *tx = *ptr;
        tx->error_code = ECONNABORTED;
        tx->complete = 1;
        pthread_cond_signal(&tx->cond);
    }

    close(conn->sockfd);
    conn->sockfd = -1;
    pthread_mutex_unlock(&conn->mutex);
}

static inline bool send_option(int sockfd, nbd_opt_t option, uint8_t *buffer, uint32_t length)
{
    uint8_t packet[16];
    *(uint64_t*)&packet[0]  = htobe64(NBD_NEWSTYLE_MAGIC);
    *(uint32_t*)&packet[8]  = htobe32(option);
    *(uint32_t*)&packet[12] = htobe32(length);

    if (write(sockfd, packet, sizeof(packet)) != sizeof(packet)) return false;
    if (length && write(sockfd, buffer, length) != length) return false;
    return true;
}

static bool init_connection(nbd_conn_t *conn, nbd_image_t *image)
{
    uint64_t magic;
    uint16_t handshake_flags;

    if (!read_uint64(conn->sockfd, &magic) || magic != NBD_MAGIC) return false;
    if (!read_uint64(conn->sockfd, &magic)) return false;

    if (magic == NBD_OLDSTYLE_MAGIC)
    {
        /* Old-style negotiation */
        TRACE("Using old-style negotiation\n");

        if (!read_uint64(conn->sockfd, &image->image_size)) return false;
        if (!read_uint16(conn->sockfd, &handshake_flags)) return false;
        if (!read_uint16(conn->sockfd, &conn->flags)) return false;
        if (!(conn->flags & NBD_FLAG_HAS_FLAGS)) return false;

        char garbage[124];
        if (read(conn->sockfd, garbage, sizeof(garbage)) != sizeof(garbage)) return false;
    }
    else if (magic == NBD_NEWSTYLE_MAGIC)
    {
        /* New-style negotiation */
        TRACE("Using new-style negotiation\n");
        if (!read_uint16(conn->sockfd, &handshake_flags)) return false;

        uint32_t client_flags = 0;
        if (handshake_flags & NBD_FLAG_NO_ZEROES) client_flags |= NBD_FLAG_C_NO_ZEROES;
        if (handshake_flags & NBD_FLAG_FIXED_NEWSTYLE) client_flags |= NBD_FLAG_C_FIXED_NEWSTYLE;

        uint32_t response = htobe32(client_flags);
        if (write(conn->sockfd, &response, sizeof(response)) != sizeof(response)) return false;

        if (client_flags & NBD_FLAG_C_FIXED_NEWSTYLE)
        {
            bool success = false;
            uint8_t info_req[] = {
                0, 0, 0, 0, 0, 1,
                (NBD_INFO_EXPORT >> 8) & 0xFF, NBD_INFO_EXPORT & 0xFF,
            };

            TRACE("Sending NBD_OPT_GO...\n");
            if (!send_option(conn->sockfd, NBD_OPT_GO, info_req, sizeof(info_req))) return false;

            for (;;)
            {
                TRACE("Waiting for a reply\n");
                if (!read_uint64(conn->sockfd, &magic) || magic != NBD_OPT_REPLY_MAGIC) return false;

                uint32_t type, length;
                if (!read_uint32(conn->sockfd, &type) || type != NBD_OPT_GO) return false;
                if (!read_uint32(conn->sockfd, &type)) return false;
                if (!read_uint32(conn->sockfd, &length)) return false;

                uint8_t data[length];
                if (length && read(conn->sockfd, data, length) != length) return false;

                if (type == NBD_REP_ACK || (int32_t)type < 0)
                {
                    TRACE("Finished processing replies\n");
                    break;
                }
                else if (type == NBD_REP_INFO)
                {
                    switch (be16toh(*(uint16_t*)&data[0]))
                    {
                    case NBD_INFO_EXPORT:
                        image->image_size = be64toh(*(uint64_t*)&data[2]);
                        conn->flags = be16toh(*(uint16_t*)&data[10]);
                        TRACE("Export size %llu\n", image->image_size);
                        success = true;
                        break;

                    default:
                        TRACE("Ignoring reply information type %u\n", type);
                    }
                }
            }

            if (!success)
            {
                /* Retry with the regular new-style command */
                client_flags &= ~NBD_FLAG_C_FIXED_NEWSTYLE;
            }
        }

        if (!(client_flags & NBD_FLAG_C_FIXED_NEWSTYLE))
        {
            TRACE("Sending NBD_OPT_EXPORT_NAME...\n");
            if (!send_option(conn->sockfd, NBD_OPT_EXPORT_NAME, NULL, 0)) return false;
            if (!read_uint64(conn->sockfd, &image->image_size)) return false;
            if (!read_uint16(conn->sockfd, &conn->flags)) return false;
            if (!(conn->flags & NBD_FLAG_HAS_FLAGS)) return false;

            if (!(client_flags & NBD_FLAG_C_NO_ZEROES))
            {
                char garbage[124];
                if (read(conn->sockfd, garbage, sizeof(garbage)) != sizeof(garbage)) return false;
            }
        }
    }
    else
    {
        errno = EINVAL;
        return false;
    }

    pthread_create(&conn->thread, NULL, process_messages, conn);
    return true;
}

static bool send_request(nbd_conn_t *conn,
                         nbd_cmd_t command,
                         uint16_t flags,
                         uint64_t offset,
                         uint32_t length,
                         uint8_t *buffer,
                         nbd_transaction_t *tx)
{
    if (!nbd_is_connected(conn)) return false;

    tx->conn = conn;
    tx->command = command;
    tx->complete = 0;
    tx->reply = buffer;
    tx->length = length;
    if (pthread_cond_init(&tx->cond, NULL)) return false;

    bool success = false;
    pthread_mutex_lock(&conn->mutex);

    for (;;)
    {
        tx->id = conn->next_txid++;
        nbd_transaction_t *added_tx = hash_table_search(&conn->tx_hash_table, tx, 1);
        if (!added_tx) goto cleanup;
        if (added_tx == tx) break;
        /* Let it repeat because this ID is in use */
    }

    uint8_t packet[28];
    *((uint32_t*)&packet[0])  = htobe32(NBD_REQUEST_MAGIC);
    *((uint16_t*)&packet[4])  = htobe16(flags);
    *((uint16_t*)&packet[6])  = htobe16(command);
    *((uint64_t*)&packet[8])  = htobe64(tx->id);
    *((uint64_t*)&packet[16]) = htobe64(offset);
    *((uint32_t*)&packet[24]) = htobe32(length);

    if (write(conn->sockfd, packet, sizeof(packet)) != sizeof(packet)) goto fail;
    if (command == NBD_CMD_WRITE && write(conn->sockfd, buffer, length) != length) goto fail;

    TRACE("[%llu] Request (command: %u, flags: %u, offset: %llu, length: %lu, buffer: %p)\n",
          tx->id,
          command,
          flags,
          offset,
          length,
          buffer);
    success = true;
    goto cleanup;

fail:
    hash_table_delete(&conn->tx_hash_table, tx);

cleanup:
    pthread_mutex_unlock(&conn->mutex);
    return success;
}

static void finish_request(nbd_transaction_t *tx)
{
    nbd_conn_t *conn = tx->conn;
    pthread_mutex_lock(&conn->mutex);
    hash_table_delete(&conn->tx_hash_table, tx);
    pthread_mutex_unlock(&conn->mutex);
}

static nbd_status_t send_wait_finish(nbd_conn_t *conn,
                                     nbd_cmd_t command,
                                     uint16_t flags,
                                     uint64_t offset,
                                     uint32_t length,
                                     uint8_t *buffer)
{
    nbd_transaction_t tx;
    if (!send_request(conn, command, flags, offset, length, buffer, &tx)) return errno;

    pthread_mutex_lock(&conn->mutex);
    while (!tx.complete) pthread_cond_wait(&tx.cond, &conn->mutex);
    pthread_mutex_unlock(&conn->mutex);

    finish_request(&tx);
    return tx.error_code;
}

nbd_status_t nbd_read(nbd_conn_t *conn, uint8_t *buffer, uint64_t offset, uint32_t length)
{
    TRACE("nbd_read(%p, %p, %llu, %lu)\n", conn, buffer, offset, length);
    return send_wait_finish(conn, NBD_CMD_READ, 0, offset, length, buffer);
}

nbd_status_t nbd_write(nbd_conn_t *conn, const uint8_t *buffer, uint64_t offset, uint32_t length)
{
    TRACE("nbd_write(%p, %p, %llu, %lu)\n", conn, buffer, offset, length);
    if (conn->flags & NBD_FLAG_READ_ONLY) return EROFS;
    return send_wait_finish(conn, NBD_CMD_WRITE, 0, offset, length, (uint8_t*)buffer);
}

nbd_status_t nbd_flush(nbd_conn_t *conn)
{
    TRACE("nbd_flush(%p)\n", conn);
    if (!(conn->flags & NBD_FLAG_SEND_FLUSH)) return 0;
    return send_wait_finish(conn, NBD_CMD_FLUSH, 0, 0, 0, NULL);
}

bool nbd_connect_tcp(const char *host, const char *port, nbd_image_t *image, nbd_conn_t *conn)
{
    static const struct addrinfo hints = { .ai_family = AF_UNSPEC, .ai_socktype = SOCK_STREAM };
    TRACE("nbd_connect_tcp(\"%s\", \"%s\", %p)\n", host, port, image);

    bool success = false;
    pthread_mutex_lock(&conn->mutex);

    if (conn->sockfd != -1)
    {
        errno = EEXIST;
        goto cleanup;
    }

    struct addrinfo *info = NULL;
    if (getaddrinfo(host, port, &hints, &info)) goto cleanup;

    for (struct addrinfo *p = info; p; p = p->ai_next)
    {
        if ((conn->sockfd = socket(p->ai_family, p->ai_socktype, p->ai_protocol)) == -1) break;
        if (connect(conn->sockfd, p->ai_addr, p->ai_addrlen) == 0) break;
        close(conn->sockfd);
        conn->sockfd = -1;
    }

    freeaddrinfo(info);
    if (conn->sockfd == -1 || (success = init_connection(conn, image))) goto cleanup;

    close(conn->sockfd);
    conn->sockfd = -1;

cleanup:
    pthread_mutex_unlock(&conn->mutex);
    return success;
}

bool nbd_connect_unix(const char *path, nbd_image_t *image, nbd_conn_t *conn)
{
    TRACE("nbd_connect_unix(\"%s\", %p)\n", path, image);

    bool success = false;
    pthread_mutex_lock(&conn->mutex);

    if (conn->sockfd != -1)
    {
        errno = EEXIST;
        goto cleanup;
    }

    struct sockaddr_un addr;
    strncpy(addr.sun_path, path, sizeof(addr.sun_path) - 1);
    addr.sun_path[sizeof(addr.sun_path) - 1] = '\0';
    memset(conn, 0, sizeof(*conn));

    if ((conn->sockfd = socket(AF_UNIX, SOCK_STREAM, 0)) == -1) goto cleanup;
    if (success = (connect(conn->sockfd, (struct sockaddr*)&addr, sizeof(addr)) >= 0
                   && init_connection(conn, image))) goto cleanup;

    close(conn->sockfd);
    conn->sockfd = -1;

cleanup:
    pthread_mutex_unlock(&conn->mutex);
    return success;
}

bool nbd_is_connected(const nbd_conn_t *conn)
{
    return conn->sockfd != -1;
}

void nbd_disconnect(nbd_conn_t *conn)
{
    TRACE("nbd_disconnect(%p)\n", conn);
    nbd_transaction_t tx;
    send_request(conn, NBD_CMD_DISC, 0, 0, 0, NULL, &tx);
    pthread_join(conn->thread, NULL);
}

void nbd_init(nbd_conn_t *conn)
{
    conn->sockfd = -1;
    conn->flags  = 0;
    pthread_mutex_init(&conn->mutex, NULL);
    conn->next_txid = 0;
    hash_table_init(&conn->tx_hash_table, tx_hash, tx_compare);
}
